import os
import re
import pickle
import json
import numpy as np
from scipy.optimize import linear_sum_assignment

# 类别映射关系
category_mapping = {
    "car": 0, "forklift": 1, "tractor": 2, "truck": 3, "flatbed": 4,
    "trailer": 5, "bicycle": 6, "cyclist": 7, "pedestrian": 8,
    "traffic_cone": 9, "water_horse": 10, "guardrail": 11, "unknown": 12,
}
class_names = list(category_mapping.keys())
num_classes = len(class_names)
distance_ranges = [10, 20, 30]  # 以米为单位

def calculate_cost_matrix(boxes1, boxes2):
    m, n = boxes1.shape[0], boxes2.shape[0]
    cost = np.zeros((m, n), dtype=float)
    for i in range(m):
        for j in range(n):
            cost[i, j] = np.linalg.norm(boxes1[i][:3] - boxes2[j][:3])
    return cost

def match_3d_boxes(boxes1, boxes2):
    cost = calculate_cost_matrix(boxes1, boxes2)
    row_idx, col_idx = linear_sum_assignment(cost)
    return np.column_stack((row_idx, col_idx))

# 加载数据
gt = pickle.load(open('dataset/rcs_0625/infos_val_01sweeps_withvelo_filter_True.pkl', 'rb'))
pre = pickle.load(open('/data2/xd/CenterPoint/work_dirs/opendata_0710/prediction.pkl', 'rb'))

gt_counts = {}
metrics = {}
metrics_dist = {}  # 新增：距离相关统计

for item in gt:
    token = item['token']
    scene = token.split("-")[2]
    gt_counts.setdefault(scene, [0] * num_classes)
    metrics.setdefault(scene, [dict(TP=0, FP=0, FN=0) for _ in range(num_classes)])
    if scene not in metrics_dist:
        metrics_dist[scene] = {
            d: [dict(TP=0, FP=0, FN=0) for _ in range(num_classes)]
            for d in distance_ranges
        }
    for n in item['gt_names']:
        idx = category_mapping.get(n)
        if idx is not None:
            gt_counts[scene][idx] += 1

# TP/FP/FN 核心逻辑
for item in gt:
    token = item['token']
    scene = token.split("-")[2]

    gt_boxes = item['gt_boxes']
    gt_labels = [category_mapping[n] for n in item['gt_names']]
    if gt_boxes.shape[0] != len(gt_labels): continue
    if token not in pre: continue

    p = pre[token]
    pre_boxes = np.array(p['box3d_lidar'])
    pre_labels = list(map(int, p['label_preds']))

    if pre_boxes.shape[0] == 0:
        for cls in gt_labels:
            metrics[scene][cls]['FN'] += 1
        continue
    if gt_boxes.shape[0] == 0:
        for cls in pre_labels:
            metrics[scene][cls]['FP'] += 1
        continue

    matches = match_3d_boxes(gt_boxes, pre_boxes)
    matched_gt = set(matches[:, 0])
    matched_pre = set(matches[:, 1])

    # 普通统计
    for i in range(len(pre_boxes)):
        if i not in matched_pre:
            metrics[scene][pre_labels[i]]['FP'] += 1
    for i in range(len(gt_boxes)):
        if i not in matched_gt:
            metrics[scene][gt_labels[i]]['FN'] += 1
    for gt_i, pr_i in matches:
        cls_gt = gt_labels[gt_i]
        cls_pr = pre_labels[pr_i]
        dist = np.linalg.norm(gt_boxes[gt_i][:3] - pre_boxes[pr_i][:3])
        if dist <= 4 and cls_gt == cls_pr:
            metrics[scene][cls_gt]['TP'] += 1
        else:
            metrics[scene][cls_pr]['FP'] += 1
            metrics[scene][cls_gt]['FN'] += 1

    # 按距离统计
    for d in distance_ranges:
        # FN
        for i, box in enumerate(gt_boxes):
            if np.linalg.norm(box[:3]) <= d and i not in matched_gt:
                cls = gt_labels[i]
                metrics_dist[scene][d][cls]['FN'] += 1
        # FP
        for i, box in enumerate(pre_boxes):
            if np.linalg.norm(box[:3]) <= d and i not in matched_pre:
                cls = pre_labels[i]
                metrics_dist[scene][d][cls]['FP'] += 1
        # TP / 误匹配
        for gt_i, pr_i in matches:
            gt_dist = np.linalg.norm(gt_boxes[gt_i][:3])
            pr_dist = np.linalg.norm(pre_boxes[pr_i][:3])
            if gt_dist <= d and pr_dist <= d:
                cls_gt = gt_labels[gt_i]
                cls_pr = pre_labels[pr_i]
                dist = np.linalg.norm(gt_boxes[gt_i][:3] - pre_boxes[pr_i][:3])
                if dist <= 4 and cls_gt == cls_pr:
                    metrics_dist[scene][d][cls_gt]['TP'] += 1
                else:
                    metrics_dist[scene][d][cls_pr]['FP'] += 1
                    metrics_dist[scene][d][cls_gt]['FN'] += 1

# 输出原始结果
final = {}
global_stats = [dict(GT=0, TP=0, FP=0, FN=0) for _ in range(num_classes)]

for scene, cm_list in metrics.items():
    print(f"\nResults for {scene}:")
    print("Class    GT    TP    FP    FN    Precision    Recall")
    scene_res = {}
    for idx in range(num_classes):
        # 取30m范围的结果
        dist_stat = metrics_dist[scene][30][idx]

        # 计算该scene中该类别在30m以内的GT数量
        GT = 0
        for item in gt:
            token = item['token']
            scene_cur = token.split("-")[2]
            if scene_cur != scene:
                continue
            gt_boxes = item['gt_boxes']
            gt_names = item['gt_names']
            for i, name in enumerate(gt_names):
                if category_mapping.get(name) == idx:
                    if np.linalg.norm(gt_boxes[i][:3]) <= 30:
                        GT += 1


        TP, FP, FN = dist_stat['TP'], dist_stat['FP'], dist_stat['FN']
        prec = TP / (TP + FP + 1e-6)
        rec = TP / (TP + FN + 1e-6)
        print(f"{class_names[idx]:<8} {GT:<5} {TP:<5} {FP:<5} {FN:<5} {prec:.3f} {rec:.3f}")
        scene_res[class_names[idx]] = {
            'GT': GT, 'TP': TP, 'FP': FP, 'FN': FN,
            'Precision': round(prec, 3), 'Recall': round(rec, 3)
        }

        # 全局统计仍然保留完整数据
        global_stats[idx]['GT'] += gt_counts[scene][idx]
        global_stats[idx]['TP'] += cm_list[idx]['TP']
        global_stats[idx]['FP'] += cm_list[idx]['FP']
        global_stats[idx]['FN'] += cm_list[idx]['FN']

    final[scene] = scene_res

# # 全局统计（完整）
# print("\nGlobal Results:")
# print("Class       GT    TP    FP    FN    Precision    Recall")
# global_res = {}
# for idx, st in enumerate(global_stats):
#     GT, TP, FP, FN = st['GT'], st['TP'], st['FP'], st['FN']
#     prec = TP / (TP + FP + 1e-6)
#     rec = TP / (TP + FN + 1e-6)
#     print(f"{class_names[idx]:<10}{GT:<6}{TP:<6}{FP:<6}{FN:<6}{prec:.3f}    {rec:.3f}")
#     global_res[class_names[idx]] = {
#         'GT': GT, 'TP': TP, 'FP': FP, 'FN': FN,
#         'Precision': round(prec, 3), 'Recall': round(rec, 3)
#     }
# final['Global'] = global_res

# # 总体结果
# total_GT = sum(st['GT'] for st in global_stats)
# total_TP = sum(st['TP'] for st in global_stats)
# total_FP = sum(st['FP'] for st in global_stats)
# total_FN = sum(st['FN'] for st in global_stats)
# overall_prec = total_TP / (total_TP + total_FP + 1e-6)
# overall_rec = total_TP / (total_TP + total_FN + 1e-6)
# print("\nOverall (all scenes & all classes):")
# print(f"GT: {total_GT}, TP: {total_TP}, FP: {total_FP}, FN: {total_FN}")
# print(f"Precision: {overall_prec:.3f}, Recall: {overall_rec:.3f}")
# final['Overall'] = {
#     'GT': total_GT, 'TP': total_TP, 'FP': total_FP, 'FN': total_FN,
#     'Precision': round(overall_prec, 3), 'Recall': round(overall_rec, 3)
# }

# 距离范围统计输出
for d in distance_ranges:
    print(f"\n===== Distance ≤ {d}m Global Stats =====")
    global_d = [dict(TP=0, FP=0, FN=0) for _ in range(num_classes)]
    for scene in metrics_dist:
        for idx, cls_stat in enumerate(metrics_dist[scene][d]):
            global_d[idx]['TP'] += cls_stat['TP']
            global_d[idx]['FP'] += cls_stat['FP']
            global_d[idx]['FN'] += cls_stat['FN']
    print("Class     TP    FP    FN    Precision    Recall")
    for idx, st in enumerate(global_d):
        TP, FP, FN = st['TP'], st['FP'], st['FN']
        prec = TP / (TP + FP + 1e-6)
        rec = TP / (TP + FN + 1e-6)
        print(f"{class_names[idx]:<10}{TP:<6}{FP:<6}{FN:<6}{prec:.3f}    {rec:.3f}")
    total_TP = sum(s['TP'] for s in global_d)
    total_FP = sum(s['FP'] for s in global_d)
    total_FN = sum(s['FN'] for s in global_d)
    overall_prec = total_TP / (total_TP + total_FP + 1e-6)
    overall_rec = total_TP / (total_TP + total_FN + 1e-6)
    print(f"Total     {total_TP:<6}{total_FP:<6}{total_FN:<6}{overall_prec:.3f}    {overall_rec:.3f}")

# 保存结果
with open('scene_class_metrics_0624.json', 'w') as f:
    json.dump(final, f, indent=2)
